# -*- coding: utf-8 -*-
import os
import shutil
import sys
import time
from collections import deque
from datetime import datetime
from pathlib import Path

from PIL import Image
from PyQt5.QtCore import QUrl, QThread, pyqtSignal, Qt, QMutex
from PyQt5.QtGui import QDesktopServices, QTextCursor, QIcon
from PyQt5.QtWidgets import QApplication, QFormLayout, QPushButton, QLineEdit, QFileDialog, QMainWindow, QWidget, \
    QLabel, QTextEdit, QVBoxLayout, QHBoxLayout
from deepface import DeepFace
from ultralytics import YOLO

'''校验是否是图片'''
def is_image(file_path):
    """检查文件是否为有效的图像文件"""
    try:
        with Image.open(file_path) as img:
            img.verify()  # 验证文件是否为有效图像而不加载整个图像到内存中
        return True
    except Exception as e:
        print(f"文件 {file_path} 无效: {e}")
        return False

'''当前目录所有图片'''
def traverse_images(directory):
    """遍历给定目录及其子目录中的所有图像文件"""
    image_extensions = ['.png', '.jpg', '.jpeg']  # 支持的图像扩展名列表
    image_files = []
    try:
        for file_path in Path(directory).glob('*'):  # 使用 glob 而不是 rglob (包含子目录)
            if file_path.suffix.lower() in image_extensions and is_image(file_path):
                image_files.append(file_path)
    except Exception as e:
            print(f"遍历目录时发生错误: {e}")
    return image_files

'''移动并且覆盖同名文件 src -- 源文件路径（字符串）dst -- 目标文件路径（字符串） '''
def move_and_overwrite(src:str, dst:str,prefix:str=""):
    # 确保源文件存在且是文件
    src_path = Path(src)
    if not src_path.is_file():
        print(f"源文件不存在或不是文件: {src}")
        return
        # 处理目标路径
    dst_path = Path(dst)
    if dst_path.is_dir():
        dst_path = dst_path / (prefix + src_path.name)
    else:
        dst_path = dst_path.parent / (prefix + dst_path.name)
    try:
        # 如果目标文件存在，则先删除它
        if dst_path.exists():
            dst_path.unlink()  # 使用 unlink 方法来删除文件
            print(f"已删除目标文件: {dst}")
        # 移动文件
        shutil.move(str(src_path), str(dst_path))
        print(f"文件已成功移动到: {dst_path}")
    except Exception as e:
        print(f"移动文件时发生错误: {e}")

# 性别映射表
gender_mapping = {
    'Man': '男性',
    'Woman': '女性'
}

'''分析图像'''
def analyze_image(image_path):
    """分析图像，获取性别和年龄信息"""
    genderList = []
    ageList = []
    try:
        # 分析图像，获取性别和年龄信息  enforce_detection=False  # 如果没有检测到人脸，仍然返回结果
        result = DeepFace.analyze(
            img_path=str(image_path),
            actions=['gender', 'age'],
            enforce_detection=False
        )
        if isinstance(result, list):
            for res in result:
                genderList.append(gender_mapping.get(res.get('dominant_gender')))  # 获取主要预测的性别
                ageList.append(res.get('age'))  # 获取预测的年龄
        elif isinstance(result, dict):
            genderList.append(gender_mapping.get(result.get('dominant_gender')))
            ageList.append(result.get('age'))
    except Exception as e:
        print(f"分析 {image_path} 时出错: {e}")
        return False, genderList, ageList
    return True,genderList, ageList

class AiYOLOThread(QThread):
    ai_signal = pyqtSignal(str,str)
    mutex = QMutex()

    def __init__(self,paths_dict):
        super().__init__()
        self._is_running = False
        self.paths_dict = paths_dict
        self.model = None

    def run(self):
            self._is_running = True
            while self._is_running:
                try:
                    '''加载预训练模型'''
                    if self.model is None:
                        self.model = YOLO("yolo11n.pt")  # 延迟加载模型

                    if self.paths_dict:
                        待识别路径 = self.paths_dict["待识别路径"]
                        image_files = traverse_images(待识别路径)
                        self.ai_signal.emit(f"AiYOLOThread: 开始检测 {len(image_files)}", "blue")
                        if image_files:
                            for img in image_files:
                                if not self._is_running:
                                    break
                                '''对图像执行推理'''
                                results = self.model(img)
                                '''COCO 数据集中类别 0 是 'person'''
                                person_detected = any(int(box.cls[0]) == 0 for result in results for box in result.boxes.cpu().numpy())
                                if person_detected:
                                    人体路径 = self.paths_dict["人体路径"]
                                    self.ai_signal.emit(f"AiYOLOThread: {str(img.absolute())} - 检测到人体 - 移到 {人体路径}", "red")
                                    move_and_overwrite(str(img.absolute()), 人体路径)
                                else:
                                    无人路径 = self.paths_dict["无人路径"]
                                    self.ai_signal.emit(f"AiAnalyzeThread: {str(img.absolute())} - 未检测到人体- 移到 {无人路径}", "Purple")
                                    move_and_overwrite(str(img.absolute()), 无人路径)
                except Exception as e:
                    self.ai_signal.emit(f"AiYOLOThread: 线程错误 {e}", "red")
                self.ai_signal.emit(f"AiYOLOThread: 休息3秒继续", "blue")
                time.sleep(3)

    def request_stop(self):
        """请求线程停止"""
        self.mutex.lock()
        self._is_running = False
        self.mutex.unlock()

    def stop(self):
        self.request_stop()
        self.quit()
        self.wait()


class AiAnalyzeThread(QThread):
    ai_signal = pyqtSignal(str,str)
    mutex = QMutex()

    def __init__(self,paths_dict):
        super().__init__()
        self._is_running = False
        self.paths_dict = paths_dict

    def run(self):
        self._is_running = True
        while self._is_running:
            if self.paths_dict:
                人体路径 = self.paths_dict["人体路径"]
                image_files = traverse_images(人体路径)
                self.ai_signal.emit(f"AiAnalyzeThread: 开始检测 {len(image_files)}", "blue")
                if image_files:
                    for img in image_files:
                        if not self._is_running:
                            break
                        '''分析图像性别和年龄'''
                        result,genderList, ageList = analyze_image(img.absolute())
                        if result:
                            if genderList and ageList:
                                if len(genderList) == 1:
                                    if genderList[0] == '男性':
                                        self.ai_signal.emit(f"AiAnalyzeThread: {str(img.absolute())} -- {genderList[0]} -- 移动 {self.paths_dict['男性路径']}", "red")
                                        move_and_overwrite(str(img.absolute()), self.paths_dict["男性路径"], prefix=f'{ageList[0]}岁_')
                                    else:
                                        self.ai_signal.emit(f"AiAnalyzeThread: {str(img.absolute())} -- {genderList[0]} -- 移动 {self.paths_dict['女性路径']}", "red")
                                        move_and_overwrite(str(img.absolute()), self.paths_dict["女性路径"], prefix=f'{ageList[0]}岁_')
                                else:
                                    self.ai_signal.emit(f"AiAnalyzeThread: 检测到 {genderList} 预测年龄 {ageList} 岁 -- 移动 {self.paths_dict['多人路径']}", "red")
                                    move_and_overwrite(str(img.absolute()), self.paths_dict["多人路径"])
                            else:
                                self.ai_signal.emit(f"AiAnalyzeThread: {str(img.absolute())}: -- 无法检测识别 -- 移动 {self.paths_dict['人妖路径']}", "Purple")
                                move_and_overwrite(str(img.absolute()), self.paths_dict["人妖路径"])
            self.ai_signal.emit(f"AiAnalyzeThread: 休息3秒继续", "blue")
            time.sleep(3)

    def request_stop(self):
        """请求线程停止"""
        self.mutex.lock()
        self._is_running = False
        self.mutex.unlock()

    def stop(self):
        self.request_stop()
        self.quit()
        self.wait()

class FileCounterThread(QThread):
    file_count_signal = pyqtSignal(str, int)

    def __init__(self, paths_dict):
        super().__init__()
        self.paths_dict = paths_dict

    def run(self):
        while (True):
            if self.paths_dict:
                # 遍历路径字典，统计每个路径下的文件数量并发送信号
                for label, path in self.paths_dict.items():
                    count = self.countFilesInDirectory(path)
                    print(f"{label} - {count}")
                    self.file_count_signal.emit(label, count)
            time.sleep(10)

    @staticmethod
    def countFilesInDirectory(directory):
        try:
            # 确保路径是绝对路径
            abs_directory = os.path.abspath(directory)
            # 使用生成器表达式以减少内存占用，并过滤文件类型
            return sum(1 for entry in os.scandir(abs_directory) if
                       entry.is_file() and entry.name.lower().endswith(('.png', '.jpg', '.jpeg')))
        except (FileNotFoundError, PermissionError) as e:
            print(f"Error: {e}")
            return 0


class DirectoryDialogDemo(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('TG: tiantian9526 【安卓群控定制化、浏览器群控定制化、服务器页面定制化】')
        self.setWindowIcon(QIcon('favicon.ico'))
        self.initUI()
        '''默认最大化'''
        self.setMinimumSize(1000, 800)
        self.show()

    def initUI(self):
        # 创建一个中心小部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        '''数据初始化'''
        self.log_buffer = deque(maxlen=1000)  # 初始化日志缓冲区，最大长度为1000条
        self.log_counter = 1  # 初始化日志序号

        self.is_running_file_counter = False  #数值统计
        self.fileCounterthread = None #数值统计线程

        self.is_ai_running = False # AI识别
        self.aithread = None # AI识别线程

        '''文件目录地址'''
        self.paths_dict = {}

        # 创建主布局
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(10, 10, 10, 10)  # 设置布局的边距
        main_layout.setSpacing(5)  # 设置控件之间的间距

        # 创建表单布局
        form_layout = QFormLayout()
        form_layout.setContentsMargins(0, 0, 0, 0)  # 设置布局的边距
        form_layout.setSpacing(5)  # 设置控件之间的间距

        # 添加标题到表单布局
        title_label = QLabel("大模型AI-识别图片、人体、性别、年龄")
        font = title_label.font()
        font.setBold(True)
        font.setPointSize(12)  # 根据需要调整字体大小
        title_label.setFont(font)
        title_label.setAlignment(Qt.AlignCenter)  # 居中显示标题
        form_layout.addRow(title_label)  # 将标题添加到表单布局的第一行

        # 创建各个路径输入框及其关联的按钮和计数标签
        self.人体路径 = QLineEdit(self)
        self.人体路径.setEnabled(False)
        self.人体路径_counter = QLabel("")
        self.人体路径_button = QPushButton("打开", self)
        self.人体路径_button.clicked.connect(lambda: self.open_folder("人体路径"))
        hlayout_人体路径 = QHBoxLayout()
        hlayout_人体路径.addWidget(self.人体路径)
        hlayout_人体路径.addWidget(self.人体路径_button)
        form_layout.addRow(QLabel("人体路径:"), hlayout_人体路径)
        form_layout.addRow(self.人体路径_counter)

        self.男性路径 = QLineEdit(self)
        self.男性路径.setEnabled(False)
        self.男性路径_counter = QLabel("")
        self.男性路径_button = QPushButton("打开", self)
        self.男性路径_button.clicked.connect(lambda: self.open_folder("男性路径"))
        hlayout_男性路径 = QHBoxLayout()
        hlayout_男性路径.addWidget(self.男性路径)
        hlayout_男性路径.addWidget(self.男性路径_button)
        form_layout.addRow(QLabel("男性路径:"), hlayout_男性路径)
        form_layout.addRow(self.男性路径_counter)

        self.女性路径 = QLineEdit(self)
        self.女性路径.setEnabled(False)
        self.女性路径_counter = QLabel("")
        self.女性路径_button = QPushButton("打开", self)
        self.女性路径_button.clicked.connect(lambda: self.open_folder("女性路径"))
        hlayout_女性路径 = QHBoxLayout()
        hlayout_女性路径.addWidget(self.女性路径)
        hlayout_女性路径.addWidget(self.女性路径_button)
        form_layout.addRow(QLabel("女性路径:"), hlayout_女性路径)
        form_layout.addRow(self.女性路径_counter)

        self.多人路径 = QLineEdit(self)
        self.多人路径.setEnabled(False)
        self.多人路径_counter = QLabel("")
        self.多人路径_button = QPushButton("打开", self)
        self.多人路径_button.clicked.connect(lambda: self.open_folder("多人路径"))
        hlayout_多人路径 = QHBoxLayout()
        hlayout_多人路径.addWidget(self.多人路径)
        hlayout_多人路径.addWidget(self.多人路径_button)
        form_layout.addRow(QLabel("多人路径:"), hlayout_多人路径)
        form_layout.addRow(self.多人路径_counter)

        self.人妖路径 = QLineEdit(self)
        self.人妖路径.setEnabled(False)
        self.人妖路径_counter = QLabel("")
        self.人妖路径_button = QPushButton("打开", self)
        self.人妖路径_button.clicked.connect(lambda: self.open_folder("人妖路径"))
        hlayout_人妖路径 = QHBoxLayout()
        hlayout_人妖路径.addWidget(self.人妖路径)
        hlayout_人妖路径.addWidget(self.人妖路径_button)
        form_layout.addRow(QLabel("人妖路径:"), hlayout_人妖路径)
        form_layout.addRow(self.人妖路径_counter)

        self.无人路径 = QLineEdit(self)
        self.无人路径.setEnabled(False)
        self.无人路径_counter = QLabel("")
        self.无人路径_button = QPushButton("打开", self)
        self.无人路径_button.clicked.connect(lambda: self.open_folder("无人路径"))
        hlayout_无人路径 = QHBoxLayout()
        hlayout_无人路径.addWidget(self.无人路径)
        hlayout_无人路径.addWidget(self.无人路径_button)
        form_layout.addRow(QLabel("无人路径:"), hlayout_无人路径)
        form_layout.addRow(self.无人路径_counter)

        self.待识别路径 = QLineEdit(self)
        self.待识别路径.setEnabled(False)
        self.待识别路径_counter = QLabel("")
        self.待识别路径_button = QPushButton("打开", self)
        self.待识别路径_button.clicked.connect(lambda: self.open_folder("待识别路径"))
        hlayout_待识别路径 = QHBoxLayout()
        hlayout_待识别路径.addWidget(self.待识别路径)
        hlayout_待识别路径.addWidget(self.待识别路径_button)
        form_layout.addRow(QLabel("待识别路径:"), hlayout_待识别路径)
        form_layout.addRow(self.待识别路径_counter)

        # 创建一个按钮，点击后会弹出目录选择对话框
        self.选择_待识别路径_button = QPushButton('选择待识别路径地址', self)
        self.选择_待识别路径_button.clicked.connect(self.openDirectoryDialog)
        form_layout.addWidget(self.选择_待识别路径_button)

        # 创建开始和结束按钮
        self.start_button = QPushButton('开始任务', self)
        self.start_button.clicked.connect(self.startAiThread)
        self.start_button.setEnabled(False)  # 初始状态下禁用开始按钮

        self.end_button = QPushButton('结束任务', self)
        self.end_button.clicked.connect(self.stopAiThread)
        self.end_button.setEnabled(False)  # 初始状态下禁用结束按钮

        button_layout = QHBoxLayout()
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.end_button)
        form_layout.addRow(button_layout)

        '''禁用路径按钮'''
        self.disablePathButtons()

        # 将表单布局添加到主布局
        main_layout.addLayout(form_layout)

        # 创建操作日志输出滚动框
        self.log_text_edit = QTextEdit(self)
        self.log_text_edit.setReadOnly(True)
        self.log_text_edit.setLineWrapMode(QTextEdit.NoWrap)  # 禁止自动换行
        self.log_text_edit.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn)  # 始终显示水平滚动条
        main_layout.addWidget(self.log_text_edit)

    '''打开目录选择对话框'''

    def openDirectoryDialog(self):
        options = QFileDialog.Options()
        # 使用 getExistingDirectory 方法来获取用户选择的目录路径
        directory = QFileDialog.getExistingDirectory(self,
                                                     "选择待识别路径地址",
                                                     "",  # 起始路径
                                                     options=options)
        if directory:
            self.log_message(f"选择了待识别路径: {directory}", color="blue")
            self.待识别路径.setText(directory)  # 将选择的目录路径设置到输入框中
            self.updateFileCounters(directory)

    '''更新文件计数器 男性 man 女性 woman 多人 multi 人类 human'''

    def updateFileCounters(self, directory):
        self.paths_dict = {
            "人体路径": os.path.join(fr'{directory}', "human"),
            "男性路径": os.path.join(fr'{directory}', "man"),
            "女性路径": os.path.join(fr'{directory}', "woman"),
            "多人路径": os.path.join(fr'{directory}', "multi"),
            "人妖路径": os.path.join(fr'{directory}', "shemale"),
            "无人路径": os.path.join(fr'{directory}', "none"),
            "待识别路径": fr'{directory}'
        }
        for label, path in self.paths_dict.items():
            if not os.path.exists(path):
                os.makedirs(path)
                self.log_message(f"创建了目录: {path}")
            line_edit_path = getattr(self, f"{label}")
            line_edit_path.setText(path)

        '''启用路径按钮'''
        self.enablePathButtons()
        '''开始统计文件数量线程'''
        self.startCounting()


    '''记录日志消息'''

    def log_message(self, message, color="black"):
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        formatted_message = f"【{current_time}】<span style='color:{color};'>{self.log_counter}. {message}</span><br>"
        self.log_buffer.append(formatted_message)
        self.log_counter += 1  # 自增序号
        self.update_log_display()  # 更新日志显示

    '''更新日志显示'''

    def update_log_display(self):
        log_content = "".join(self.log_buffer)  # 将缓冲区中的日志消息合并成一个字符串
        self.log_text_edit.clear()  # 清空现有内容
        self.log_text_edit.insertHtml(log_content)  # 插入 HTML 内容
        cursor = self.log_text_edit.textCursor()  # 移动光标到底部
        cursor.movePosition(QTextCursor.End)
        self.log_text_edit.setTextCursor(cursor)

    '''打开指定目录'''

    def open_folder(self, label):
        line_edit_path = getattr(self, label)
        folder_path = line_edit_path.text()
        if folder_path and os.path.exists(folder_path):
            QDesktopServices.openUrl(QUrl.fromLocalFile(folder_path))
            self.log_message(f"打开了文件夹: {folder_path}", color="blue")

    '''开始统计文件数量线程'''

    def startCounting(self):
        if self.paths_dict:
            if self.is_running_file_counter:
                self.log_message("统计已经在运行中", color="orange")
                if self.fileCounterthread:
                    self.fileCounterthread.paths_dict = self.paths_dict
                    self.log_message(f"切换路口:{self.paths_dict['待识别路径']}",color="blue")
                return

            if not self.is_running_file_counter:
                self.is_running_file_counter = True
                self.log_message("开始统计文件数量线程", color="blue")
                self.fileCounterthread = FileCounterThread(self.paths_dict)
                self.fileCounterthread.file_count_signal.connect(self.updateCounterLabels)
                self.fileCounterthread.start()
                return

    '''开启ai线程'''

    def startAiThread(self):
        if not self.is_ai_running:
            self.log_message("开启线程", color="pink")
            self.start_button.setEnabled(False)
            self.end_button.setEnabled(True)
            self.选择_待识别路径_button.setEnabled(False)
            '''------------------------------------------------------------'''
            self.ai_YOLO_Thread = AiYOLOThread(paths_dict=self.paths_dict)
            self.ai_YOLO_Thread.ai_signal.connect(self.log_message)
            self.ai_YOLO_Thread.finished.connect(lambda: self.onAiThreadFinished("ai_YOLO_Thread"))
            '''------------------------------------------------------------'''
            self.ai_analyze_Thread = AiAnalyzeThread(paths_dict=self.paths_dict)
            self.ai_analyze_Thread.ai_signal.connect(self.log_message)
            self.ai_analyze_Thread.finished.connect(lambda: self.onAiThreadFinished("ai_analyze_Thread"))

            self.ai_YOLO_Thread.start()
            self.ai_analyze_Thread.start()
            self.is_ai_running = True

    '''停止AI线程'''

    def stopAiThread(self):
        if self.ai_YOLO_Thread is not None and self.is_ai_running:
            self.ai_YOLO_Thread.request_stop()  # 调用线程的停止方法而不是直接终止

        if self.ai_analyze_Thread is not None and self.is_ai_running:
            self.ai_analyze_Thread.request_stop()  # 调用线程的停止方法而不是直接终止

    '''停止线程'''
    def onAiThreadFinished(self,threadName):
        self.log_message(f"{threadName}线程已结束", color="pink")
        if threadName:
            if threadName == 'ai_YOLO_Thread':
                self.ai_YOLO_Thread = None  # 清除对线程的引用
            elif threadName == 'ai_analyze_Thread':
                self.ai_analyze_Thread = None  # 清除对线程的引用
        if not self.ai_YOLO_Thread and not self.ai_analyze_Thread:
            self.is_ai_running = False
            self.start_button.setEnabled(True)
            self.end_button.setEnabled(False)
            self.选择_待识别路径_button.setEnabled(True)
            self.is_ai_running = False

    '''更新计数标签'''

    def updateCounterLabels(self, label, count):
        counter_label = getattr(self, f"{label}_counter")
        counter_label.setText(f"文件数量: {count}")
        if count == 0:
            counter_label.setStyleSheet("color: green;")
            self.log_message(f"{label} 文件数量:{count}", color="green")
        else:
            counter_label.setStyleSheet("color: red;")
            self.log_message(f"{label} 文件数量:{count}", color="red")

    '''禁用路径按钮'''

    def disablePathButtons(self):
        self.人体路径_button.setEnabled(False)
        self.男性路径_button.setEnabled(False)
        self.女性路径_button.setEnabled(False)
        self.多人路径_button.setEnabled(False)
        self.人妖路径_button.setEnabled(False)
        self.无人路径_button.setEnabled(False)
        self.待识别路径_button.setEnabled(False)
        self.start_button.setEnabled(False)  # 禁用开始按钮
        self.end_button.setEnabled(False)  # 禁用结束按钮

    '''启用路径按钮'''

    def enablePathButtons(self):
        self.人体路径_button.setEnabled(True)
        self.男性路径_button.setEnabled(True)
        self.女性路径_button.setEnabled(True)
        self.多人路径_button.setEnabled(True)
        self.人妖路径_button.setEnabled(True)
        self.无人路径_button.setEnabled(True)
        self.待识别路径_button.setEnabled(True)
        self.start_button.setEnabled(True)  # 启用开始按钮
        self.end_button.setEnabled(False)  # 禁用结束按钮


if __name__ == '__main__':
    app = QApplication(sys.argv)
    mainWin = DirectoryDialogDemo()
    sys.exit(app.exec_())
